"""
基于 LangChain 的数据分析智能体
"""
from langchain_core.prompts import ChatPromptTemplate, SystemMessagePromptTemplate, HumanMessagePromptTemplate
from langchain_core.output_parsers import StrOutputParser, PydanticOutputParser
from langchain_core.runnables import RunnablePassthrough, RunnableLambda
from agents.base_agent import StockAgent
from pydantic import BaseModel, Field
from typing import List, Dict, Any
from colorama import Fore, Style
import json


class StockAnalysisResult(BaseModel):
    """股票分析结果模型"""
    technical_analysis: str = Field(description="技术分析结果")
    fundamental_analysis: str = Field(description="基本面分析结果")
    risk_assessment: str = Field(description="风险评估")
    market_sentiment: str = Field(description="市场情绪分析")
    key_insights: List[str] = Field(description="关键洞察点")
    summary: str = Field(description="分析总结")


class DataAnalystAgent(StockAgent):
    """基于 LangChain 的数据分析智能体"""
    
    def __init__(self):
        super().__init__("data_analyst")
        self.analysis_chain = self._create_analysis_chain()
        self.output_parser = PydanticOutputParser(pydantic_object=StockAnalysisResult)

    def _create_system_prompt(self) -> str:
        """重写系统提示词"""
        return """你是一个专业的股票数据分析师，具有丰富的金融市场分析经验。

你的分析任务包括：
1. 技术分析：分析价格走势、成交量、技术指标等
2. 基本面分析：评估公司财务状况、行业地位等
3. 风险评估：识别潜在风险因素
4. 市场情绪：分析新闻和市场反应
5. 关键洞察：提炼最重要的分析要点

请基于提供的数据进行深入、客观的分析，避免主观臆断。"""

    def _create_analysis_chain(self):
        """创建分析处理链"""
        
        # 创建分析提示词模板
        analysis_prompt = ChatPromptTemplate.from_messages([
            SystemMessagePromptTemplate.from_template(self._create_system_prompt()),
            HumanMessagePromptTemplate.from_template("""
请分析以下股票数据报告：

{raw_data_report}

请从以下几个维度进行分析：
1. 技术分析：基于价格数据分析趋势和模式
2. 基本面分析：评估公司基本情况
3. 风险评估：识别主要风险因素
4. 市场情绪：基于新闻和市场反应分析情绪
5. 关键洞察：总结最重要的发现
6. 分析总结：给出综合性结论

请提供详细、专业的分析报告。
""")
        ])
        
        def enhance_analysis(analysis_result: str) -> str:
            """增强分析结果"""
            enhanced_template = """
# 股票数据深度分析报告

{analysis_result}

---
**分析方法说明：**
- 使用 LangChain 智能分析链
- 结合技术分析和基本面分析
- 基于历史数据和市场信息
- 采用多维度风险评估模型

**免责声明：**
本分析报告仅供参考，不构成投资建议。投资有风险，决策需谨慎。
"""
            return enhanced_template.format(analysis_result=analysis_result)
        
        # 创建处理链
        chain = (
            {"raw_data_report": RunnablePassthrough()}
            | analysis_prompt
            | self.llm
            | StrOutputParser()
            | RunnableLambda(enhance_analysis)
        )
        
        return chain

    def _create_structured_analysis_chain(self):
        """创建结构化分析链（可选的高级功能）"""
        
        structured_prompt = ChatPromptTemplate.from_messages([
            SystemMessagePromptTemplate.from_template(
                self._create_system_prompt() + "\n\n" + 
                "请按照以下JSON格式返回分析结果：\n{format_instructions}"
            ),
            HumanMessagePromptTemplate.from_template("""
分析以下数据：
{raw_data_report}
""")
        ])
        
        chain = (
            {
                "raw_data_report": RunnablePassthrough(),
                "format_instructions": lambda _: self.output_parser.get_format_instructions()
            }
            | structured_prompt
            | self.llm
            | self.output_parser
        )
        
        return chain

    def execute(self, raw_data_report: str) -> str:
        """执行数据分析任务"""
        try:
            print(Fore.MAGENTA + f"\n    - [LangChain] 启动数据分析链...")
            
            # 使用 LangChain 分析链处理
            analysis_result = self.analysis_chain.invoke(raw_data_report)
            
            print(Fore.GREEN + "    - [LangChain] 数据分析链处理完成")
            return analysis_result
            
        except Exception as e:
            error_msg = f"数据分析过程中发生错误: {str(e)}"
            print(Fore.RED + f"    - [Error] {error_msg}")
            return error_msg

    def execute_structured_analysis(self, raw_data_report: str) -> Dict[str, Any]:
        """执行结构化数据分析（返回结构化结果）"""
        try:
            print(Fore.MAGENTA + f"\n    - [LangChain] 启动结构化分析链...")
            
            structured_chain = self._create_structured_analysis_chain()
            result = structured_chain.invoke(raw_data_report)
            
            print(Fore.GREEN + "    - [LangChain] 结构化分析完成")
            return result.dict()
            
        except Exception as e:
            error_msg = f"结构化分析过程中发生错误: {str(e)}"
            print(Fore.RED + f"    - [Error] {error_msg}")
            return {"error": error_msg}

